# File   : augment.py
# Author : Zhengkun Tian
# Email  : zhengkun.tian@outlook.com

import oneflow as torch
import oneflow.nn as nn
from oasr.frontend import BaseFrontEnd


class SpecAugmentFrontEnd(BaseFrontEnd):
    def __init__(self, freq_mask_rate=0.3, freq_mask_num=2, time_mask_rate=0.05, time_mask_num=2):
        super(SpecAugmentFrontEnd, self).__init__()

        self.freq_mask_rate = freq_mask_rate
        self.freq_mask_num = freq_mask_num

        self.time_mask_rate = time_mask_rate
        self.time_mask_num = time_mask_num

    def forward(self, inputs, inputs_mask):


        return inputs, inputs_mask

    def inference(self, inputs, inputs_mask):
        return inputs, inputs_mask
